from torch.utils.data import DataLoader, Dataset
import numpy as np
import glob
from libtiff import TIFF
import itertools
from skimage.exposure import match_histograms
import torch

class RssRai_dataset(Dataset):
    def __init__(self, root, size=512, slide=256):
        image_path_list_2018 = glob.glob(root + '/img_2018' + '/*.tif')
        image_path_list_2018.sort()
        image_path_list_2017 = glob.glob(root + '/img_2017' + '/*.tif')
        image_path_list_2017.sort()
        image_path_list_mask = glob.glob(root + '/mask' + '/*.tif')
        image_path_list_mask.sort()
        self.image_path_list_2018 = image_path_list_2018
        self.image_path_list_2017 = image_path_list_2017
        self.image_path_list_mask = image_path_list_mask

    def __getitem__(self, index):
        image1 = TIFF.open(self.image_path_list_2017[index], mode='r').read_image()[:, :, [2, 1, 0]].transpose((2, 0, 1))
        image2 = TIFF.open(self.image_path_list_2018[index], mode='r').read_image()[:, :, [2, 1, 0]].transpose((2, 0, 1))
        image2 = match_histograms(image2, image1, multichannel=True)
        label = TIFF.open(self.image_path_list_mask[index], mode='r').read_image()
        label = (label / 255).astype(np.uint8)
        return torch.from_numpy(image1.astype(np.float32)), torch.from_numpy(image2.astype(np.float32)), torch.from_numpy(label).float().unsqueeze_(0)

    def __len__(self):
        return len(self.image_path_list_2018)

if __name__ == '__main__':
    test_loader = DataLoader(RssRai_dataset('F:/rssrai2019_change_detection/train/train'), batch_size=3, shuffle=True, num_workers=1)
    test_iter = iter(test_loader)
    x1, x2, target = next(test_iter)
    print(x1.shape, x2.shape, target.shape)